import json
import asyncio
import traceback
from typing import List

from uvicorn import Config, Server
from fastapi import FastAPI, WebSocket

from config import HOST, PORT
from handler import bilibili
from logger import logger

app = FastAPI()


class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.active_connections.append(ws)
        logger.info(f"已连接，总数：{len(self.active_connections):>2}")

    def disconnect(self, ws: WebSocket):
        self.active_connections.remove(ws)
        logger.info(f"已断开，总数：{len(self.active_connections):>2}")

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


manager = ConnectionManager()
queue = asyncio.Queue(1000)


@app.websocket("/endpoint")
async def endpoint(websocket: WebSocket):
    await manager.connect(websocket)

    try:
        while True:
            data = await queue.get()
            await manager.broadcast(json.dumps(data))

    except:
        logger.error(traceback.format_exc())
        manager.disconnect(websocket)


async def main():
    logger.success("启动中")
    await asyncio.gather(
        Server(config=Config(app, host=HOST, port=PORT, log_level="critical")).serve(),
        bilibili(queue),
    )


if __name__ == "__main__":
    asyncio.run(main())
